feat: JIT-friendly Einstein radius helper + AnalysisDataset jit latent mode - #435
Merged
Merged
Conversation
…t mode
Two coupled changes that unlock JAX-mode latent computation for
lensing pipelines (Euclid, clusters, anywhere `compute_latent_samples`
runs with `use_jax=True`):
1. LensCalc.einstein_radius_jit_from(init_guess, ...) — new ~95-line
helper that bypasses the non-jit-able pieces of the existing
`einstein_radius_via_zero_contour_from`:
- Requires a static `init_guess` (caller provides a small fan of
JAX-array seeds) — skips the marching-squares seed search in
`_init_guess_from_coarse_grid` which uses skimage and is not
JAX-traceable.
- Skips `ZeroSolver.path_reduce`, which is explicitly documented
as un-jit-able (variable-length output).
- Computes the enclosed area on the raw NaN-padded paths array
via JAX-vectorised shoelace with `jnp.where(jnp.isfinite, ...)`
masking. Takes max area across seeds (robust to duplicate seeds
landing on the same curve).
- Returns a scalar `jax.Array` directly (no `float()` cast).
- Reuses the `(f, ZeroSolver)` cache from PR #434 so the warm
call inside a jit-cached `compute_latent_variables` hits the
JAX compile cache.
2. AnalysisDataset.LATENT_BATCH_MODE = "jit" — overrides PyAutoFit's
default "vmap". Required because the new helper (and
ZeroSolver in general) is not vmap-safe: upstream
`jax_zero_contour.zero_contour_finder` uses `jax.lax.cond` and
`jax.lax.while_loop` for early termination which combine poorly
with vmap. AnalysisImaging, AnalysisInterferometer, and the
PyAutoLens variants inherit this override automatically.
Companion PR: rhayes777/PyAutoFit#<TBD> adds the LATENT_BATCH_MODE
attribute that this override targets. Workspace PR:
PyAutoLabs/euclid_strong_lens_modeling_pipeline#14 wires the
Euclid effective_einstein_radius latent through the new helper.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two coupled changes that unlock JAX-mode latent computation for lensing pipelines:
LensCalc.einstein_radius_jit_from(init_guess, ...)— new ~95-line helper, fully JIT-friendly. Bypasses the non-jit-able pieces of the existingeinstein_radius_via_zero_contour_from: no marching-squares seed search (_init_guess_from_coarse_griduses skimage and breaks JAX trace), noZeroSolver.path_reduce(upstream documents it as un-jit-able), nofloat()cast. Computes the enclosed area on the raw NaN-paddedpathsarray via JAX-vectorised shoelace withjnp.where(jnp.isfinite, ...)masking, takes max area across seeds (robust to duplicate seeds on the same curve), returns a scalarjax.Array. Reuses the(f, ZeroSolver)cache from perf(lens_calc): cache (f, ZeroSolver) for zero_contour critical curves #434.AnalysisDataset.LATENT_BATCH_MODE = "jit"— overrides PyAutoFit's default"vmap". Required because both the new helper and the underlyingZeroSolverusejax.lax.cond/jax.lax.while_loopfor early termination, which upstream explicitly documents as vmap-incompatible. All PyAutoGalaxy/PyAutoLens analysis classes inherit this override automatically.Together these make
compute_latent_sampleswithuse_jax=Truework for lensing latents — previously impossible because the only Einstein-radius computation available waseinstein_radius_via_zero_contour_from(Python loops + variable-length lists, not jit-able) oreinstein_radius_from(grid=...)(marching squares + not JAX-traceable).API Changes
Two additions to PyAutoGalaxy: a new public method
LensCalc.einstein_radius_jit_from, and a class attribute overrideAnalysisDataset.LATENT_BATCH_MODE = "jit". No existing API surface is modified or removed.See full details below.
Test Plan
LensCalc.einstein_radius_jit_fromexists with requiredinit_guessargument and the expected kwargs (delta,N,pixel_scales,tol,max_newton). JAX-runtime behaviour exercised in the workspace integration (no JAX in library unit tests per project policy).AnalysisDataset.LATENT_BATCH_MODE == "jit".pytest test_autogalaxy/operate/ test_autogalaxy/analysis/— 78 passed locally.euclid_strong_lens_modeling_pipelinestart_here.py underPYAUTO_TEST_MODE=1): latent.effective_einstein_radius = 2.1002 arcsec on the prior-median MGE tracer. ~480 ms/sample on CPU after the ~10 s first-sample compile.PYAUTO_DISABLE_JAX=1): 6/6 PASS — no regression on the legacy path.Full API Changes (for automation & release notes)
Added
LensCalc.einstein_radius_jit_from(init_guess, delta=0.05, N=500, pixel_scales=(0.05, 0.05), tol=1e-6, max_newton=5) -> jax.Array— JIT-friendly Einstein radius via zero-contour. Takes a staticinit_guessof shape(n_seeds, 2), returns a scalarjax.Array. Designed for use insidejax.jit(...)per-sample; do NOT combine withjax.vmap(upstreamjax_zero_contour.ZeroSolveruseslax.cond/lax.while_loopnot safe under vmap).AnalysisDataset.LATENT_BATCH_MODE: str = "jit"— class attribute override of PyAutoFit's"vmap"default. Inherited byAnalysisImaging,AnalysisInterferometer,AnalysisEllipse,AnalysisQuantity, and downstream PyAutoLens subclasses.Changed Signature
Changed Behaviour
compute_latent_sampleswithuse_jax=Truenow dispatches via per-samplejax.jitinstead ofjax.jit(jax.vmap(...)). JIT compile cache is reused across samples; speed is faster than per-sample NumPy but slower than the (vmap-incompatible) batched ideal. Required to unlock JAX-mode lensing latents at all.Migration
einstein_radius_via_zero_contour_fromare unchanged. Neweinstein_radius_jit_fromis an additive opt-in for JIT contexts.Companion PRs
https://github.com/rhayes777/PyAutoFit/pull/1288— adds theLATENT_BATCH_MODEattribute machinery this override depends on.https://github.com/PyAutoLabs/euclid_strong_lens_modeling_pipeline/pull/15— wires the Euclideffective_einstein_radiuslatent througheinstein_radius_jit_from.Closes PyAutoLabs/euclid_strong_lens_modeling_pipeline#14 alongside the companion PRs.
🤖 Generated with Claude Code